#pragma once
#ifndef OBR_TEXTUREDS_H
#define OBR_TEXTUREDS_H

#include "DScommon.h"

namespace OBR
{
enum class TextureSampleMode
{
    CLOSEST,
    XLINEAR
};

template <typename DataType>
__forceinline__ __device__ DataType interp_1D(DataType &d0, DataType &d1, float x)
{
    return d0 * (1.0f - x) + d1 * x;
}

template <typename DataType>
__forceinline__ __device__ DataType interp_2D(DataType &d00, DataType &d01, DataType &d10,
                                              DataType &d11, float x, float y)
{
    DataType mix0 = interp_1D<DataType>(d00, d01, x);
    DataType mix1 = interp_1D<DataType>(d10, d11, x);
    return interp_1D<DataType>(mix0, mix1, y);
}

template <typename DataType>
__forceinline__ __device__ DataType interp_3D(DataType &d000, DataType &d001, DataType &d010,
                                              DataType &d011, DataType &d100, DataType &d101,
                                              DataType &d110, DataType &d111, float x, float y,
                                              float z)
{
    DataType mix0 = interp_2D<DataType>(d000, d001, d010, d011, x, y);
    DataType mix1 = interp_2D<DataType>(d100, d101, d110, d111, x, y);
    return interp_1D<DataType>(mix0, mix1, z);
}

template <typename DataType> struct DeviceTexture2D
{
    DataType *data;
    int2 size;
    __forceinline__ __device__ DataType &at(int x, int y)
    {
        return data[y * this->size.x + x];
    }
    __forceinline__ __device__ DataType sample(float2 &uv,
                                               TextureSampleMode mode = TextureSampleMode::CLOSEST)
    {
        if (mode == TextureSampleMode::CLOSEST)
        {
            int x = clamp(roundf(uv.x * this->size.x), 0.0f, this->size.x - 1.0f);
            int y = clamp(roundf(uv.y * this->size.y), 0.0f, this->size.y - 1.0f);
            return this->at(x, y);
        }
        else if (mode == TextureSampleMode::XLINEAR)
        {
            int x0 = clamp(floorf(uv.x * this->size.x), 0.0f, this->size.x - 1.0f);
            int x1 = clamp(floorf(uv.x * this->size.x) + 1, 0.0f, this->size.x - 1.0f);
            int y0 = clamp(floorf(uv.y * this->size.y), 0.0f, this->size.y - 1.0f);
            int y1 = clamp(floorf(uv.y * this->size.y) + 1, 0.0f, this->size.y - 1.0f);
            float x_interp = (clamp(uv.x * this->size.x, 0.0f, (float)this->size.x) - x0);
            float y_interp = (clamp(uv.y * this->size.y, 0.0f, (float)this->size.y) - y0);
            DataType &c00 = this->at(x0, y0);
            DataType &c01 = this->at(x1, y0);
            DataType &c10 = this->at(x0, y1);
            DataType &c11 = this->at(x1, y1);
            return interp_2D<DataType>(c00, c01, c10, c11, x_interp, y_interp);
        }
        return this->at(0, 0);
    }
};

template <typename DataType> struct DeviceTexture3D
{
    DataType *data;
    int3 size;
    __forceinline__ __device__ DataType &at(int x, int y, int z)
    {
        return data[z * this->size.y * this->size.x + y * this->size.x + x];
    }
    __forceinline__ __device__ DataType sample(float3 &pos,
                                               TextureSampleMode mode = TextureSampleMode::CLOSEST)
    {
        if (mode == TextureSampleMode::CLOSEST)
        {
            int x = clamp(roundf(pos.x * this->size.x), 0.0f, this->size.x - 1.0f);
            int y = clamp(roundf(pos.y * this->size.y), 0.0f, this->size.y - 1.0f);
            int z = clamp(roundf(pos.z * this->size.z), 0.0f, this->size.z - 1.0f);
            return this->at(x, y, z);
        }
        else if (mode == TextureSampleMode::XLINEAR)
        {
            int x0 = clamp(floorf(pos.x * this->size.x), 0.0f, this->size.x - 1.0f);
            int x1 = clamp(ceilf(pos.x * this->size.x), 0.0f, this->size.x - 1.0f);
            int y0 = clamp(floorf(pos.y * this->size.y), 0.0f, this->size.y - 1.0f);
            int y1 = clamp(ceilf(pos.y * this->size.y), 0.0f, this->size.y - 1.0f);
            int z0 = clamp(floorf(pos.z * this->size.z), 0.0f, this->size.z - 1.0f);
            int z1 = clamp(ceilf(pos.z * this->size.z), 0.0f, this->size.z - 1.0f);
            float x_interp = (clamp(pos.x * this->size.x, 0.0f, this->size.x - 1.0f) - x0);
            float y_interp = (clamp(pos.y * this->size.y, 0.0f, this->size.y - 1.0f) - y0);
            float z_interp = (clamp(pos.z * this->size.z, 0.0f, this->size.z - 1.0f) - z0);
            DataType &c000 = this->at(x0, y0, z0);
            DataType &c001 = this->at(x1, y0, z0);
            DataType &c010 = this->at(x0, y1, z0);
            DataType &c011 = this->at(x1, y1, z0);
            DataType &c100 = this->at(x0, y0, z1);
            DataType &c101 = this->at(x1, y0, z1);
            DataType &c110 = this->at(x0, y1, z1);
            DataType &c111 = this->at(x1, y1, z1);
            return interp_3D<DataType>(                         // trilinear interp
                c000, c001, c010, c011, c100, c101, c110, c111, // data
                x_interp, y_interp, z_interp                    // interp
            );
        }
        return this->at(0, 0);
    }
};
} // namespace OBR

#endif